import os
import sys
import json
from typing import Dict, Any, List

from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import PROGRAMMER_PROMPT, CLASS_PROMPT, FUNCTION_PROMPT
from utils.code_parser import extract_function_code, extract_class_code, extract_header
from agents.base_agent import BaseAgent

class ProgrammerAgent(BaseAgent):
    def __init__(self, vector_store, project_path, call_graph, file_vector):
        super().__init__("程序员")
        self.project_path = project_path
        self.call_graph = call_graph
        self.vector_store = vector_store
        self.file_vector = file_vector
        
        # 加载不同类型的提示词
        self.load_prompts()
        
        # 设置向量检索器
        self.retriever = vector_store.as_retriever(search_kwargs={"k": 5})
    
    def load_prompts(self):
        """加载各种提示词模板"""
        prompt_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "prompts")
        
        # 加载通用程序员提示词
        with open(PROGRAMMER_PROMPT, 'r', encoding='utf-8') as f:
            programmer_template = f.read()
        self.prompt = PromptTemplate(
            template=programmer_template,
            input_variables=["architect_overview", "call_graph", "code_segments", "file_path", "file_vector"]
        )
        
        # 加载类文档提示词
        with open(CLASS_PROMPT, 'r', encoding='utf-8') as f:
            class_template = f.read()
        self.class_prompt = PromptTemplate(
            template=class_template,
            input_variables=["class_name", "file_path", "class_code", "class_header","call_graph", "related_context"]
        )
        
        # 加载函数文档提示词
        with open(FUNCTION_PROMPT, 'r', encoding='utf-8') as f:
            function_template = f.read()
        self.function_prompt = PromptTemplate(
            template=function_template,
            input_variables=["function_name", "file_path", "class_info", "function_code", "call_graph", "related_context"]
        )
    
    def analyze_file_details(self, file_path, architect_overview):
        """详细分析文件并生成函数和类的详细说明 (保留旧方法以兼容)"""
        # 获取文件中的代码段
        code_segments = self._extract_relevant_code_segments(file_path)
        
        # 准备调用图信息
        file_functions = []
        for func_name, info in self.call_graph['function_info'].items():
            if info['file'] == file_path:
                file_functions.append({
                    'name': info['name'],
                    'callers': self.call_graph['caller_graph'].get(func_name, []),
                    'callees': self.call_graph['call_graph'].get(func_name, []),
                    'class': info.get('class', None)
                })
        
        # 准备输入
        inputs = {
            'architect_overview': architect_overview,
            'call_graph': str(file_functions),
            'code_segments': str(code_segments),
            'file_path': file_path,
            'file_vector': str(self.file_vector) if self.file_vector is not None else None
        }
        
        # 使用流式输出运行
        result = self.run_with_streaming(
            self.prompt,
            inputs,
            step="分析代码实现",
            file_path=file_path
        )
        
        return {
            'file_path': file_path,
            'details': result,
            'inputs': {
                'architect_overview': architect_overview,
                'call_graph': str(file_functions),
                'code_segments': str({k: v['type'] for k, v in code_segments.items()}),
                'file_path': file_path,
                'file_vector': str(self.file_vector) if self.file_vector is not None else None
            }
        }
    
    def generate_class_doc(self, file_path, class_name):
        """为单个类生成详细文档"""
        # 提取类代码
        class_code = extract_class_code(self.project_path, file_path, class_name)
        class_header = extract_header(self.project_path, file_path)
        
        # 获取相关上下文和调用关系
        related_context = self._get_relevant_context(file_path, class_name)
        call_graph = self._get_call_graph_for_class(file_path, class_name)
        
        # 生成文档
        result = self.run_with_streaming(
            self.class_prompt,
            {
                "class_name": class_name,
                "file_path": file_path,
                "class_code": class_code,
                "class_header": class_header,               
                "call_graph": json.dumps(call_graph, indent=2),
                "related_context": related_context

            },
            step=f"为类 {class_name} 生成文档",
            file_path=file_path
        )
        
        return {
            "type": "class",
            "name": class_name,
            "file_path": file_path,
            "documentation": result
        }
    
    def generate_function_doc(self, file_path, function_name, class_name=None):
        """为单个函数生成详细文档"""
        # 提取函数代码
        function_code = extract_function_code(self.project_path, file_path, function_name, class_name)
        
        # 获取相关上下文和调用关系
        related_context = self._get_relevant_context(file_path, function_name)
        call_graph = self._get_call_graph_for_function(file_path, function_name)
        
        # 准备类信息
        class_info = f"所属类: {class_name}" if class_name and class_name != "None" else ""
        
        # 生成文档
        result = self.run_with_streaming(
            self.function_prompt,
            {
                "function_name": function_name,
                "file_path": file_path,
                "class_info": class_info,
                "function_code": function_code,
                "call_graph": json.dumps(call_graph, indent=2),
                "related_context": related_context
            },
            step=f"为函数 {function_name} 生成文档",
            file_path=file_path
        )
        
        return {
            "type": "function",
            "name": function_name,
            "class_name": class_name,
            "file_path": file_path,
            "documentation": result
        }
    
    def _extract_relevant_code_segments(self, file_path):
        """提取相关的代码段，使用调用图和向量检索来获取更全面的上下文，仅供老方法使用"""
        code_segments = {}
        
        # 1. 首先从调用图中获取文件中的函数和类
        file_functions = []
        for func_name, info in self.call_graph['function_info'].items():
            if info['file'] == file_path:
                file_functions.append(func_name)
                # 如果是类的方法
                if 'class' in info and info['class']:
                    if info['class'] not in code_segments:
                        class_code = extract_class_code(self.project_path, file_path, info['class'])
                        code_segments[info['class']] = {'type': 'class', 'code': class_code}
                # 如果是普通函数
                else:
                    function_code = extract_function_code(self.project_path, file_path, info['name'], info['class'])
                    code_segments[info['name']] = {'type': 'function', 'code': function_code}
        
        # 2. 获取调用关系，找出调用当前文件函数的函数和被当前文件函数调用的函数
        related_functions = set()
        for func_name in file_functions:
            # 添加调用当前函数的函数（调用者）
            callers = self.call_graph['caller_graph'].get(func_name, [])
            related_functions.update(callers)
            
            # 添加被当前函数调用的函数（被调用者）
            callees = self.call_graph['call_graph'].get(func_name, [])
            related_functions.update(callees)
        
        # 3. 使用向量检索获取相关代码段
        # 首先构建查询字符串，包含文件名和相关函数
        query = f"文件: {file_path} 函数: {', '.join(file_functions)}"
        
        # 使用检索器获取相关文档
        relevant_docs = self.retriever.invoke(query)
        
        # 4. 将检索到的相关代码添加到代码段中
        for i, doc in enumerate(relevant_docs):
            # 提取文档内容
            content = doc.page_content
            metadata = doc.metadata
            
            # 获取源文件和可能的函数名
            source_file = metadata.get('source', '')
            
            # 如果这是一个新的相关代码段（不在当前文件中），添加到结果中
            if source_file and source_file != file_path:
                # 尝试从元数据或内容中提取函数/类名
                segment_name = metadata.get('function', f"related_segment_{i}")
                
                # 检查是否是调用关系中的函数
                if segment_name in related_functions:
                    code_segments[f"related_{segment_name}"] = {
                        'type': 'related_function',
                        'code': content,
                        'source': source_file
                    }
        
        # 5. 如果有文件向量，使用它来增强检索
        if self.file_vector is not None:
            try:
                # 使用文件向量直接搜索相似代码
                similar_docs = self.vector_store.similarity_search_by_vector(
                    self.file_vector,
                    k=3  # 获取3个最相似的结果
                )
                
                # 将相似文档添加到代码段中
                for i, doc in enumerate(similar_docs):
                    content = doc.page_content
                    metadata = doc.metadata
                    source_file = metadata.get('source', '')
                    
                    # 避免重复添加当前文件
                    if source_file and source_file != file_path:
                        segment_name = metadata.get('function', f"vector_similar_{i}")
                        if segment_name not in code_segments:
                            code_segments[f"vector_similar_{segment_name}"] = {
                                'type': 'vector_similar',
                                'code': content,
                                'source': source_file
                            }
            except Exception as e:
                print(f"使用文件向量搜索时出错: {str(e)}")
        
        print(f"共提取到 {len(code_segments)} 个代码段，包括文件内部代码和相关代码")
        return code_segments
    
    def _get_relevant_context(self, file_path, item_name):
        """获取相关上下文"""
        # 构建查询
        query = f"文件: {file_path} 项目: {item_name} 功能和作用"
        
        # 使用检索器获取相关文档
        try:
            docs = self.retriever.invoke(query)
            context = "\n\n".join([f"--- {doc.metadata.get('source', '未知文件')} ---\n{doc.page_content}" for doc in docs])
            return context
        except Exception as e:
            print(f"获取相关上下文时出错: {str(e)}")
            return "获取相关上下文失败"
    
    def _get_call_graph_for_class(self, file_path, class_name):
        """获取类的调用关系"""
        class_methods = []
        callers = {}
        callees = {}
        
        # 找出属于这个类的所有方法
        for func_name, info in self.call_graph['function_info'].items():
            if info['file'] == file_path and info.get('class') == class_name:
                method_name = info['name']
                class_methods.append(method_name)
                
                # 收集调用这个方法的函数
                method_full_name = f"{file_path}::{method_name}"
                callers[method_name] = self.call_graph['caller_graph'].get(method_full_name, [])
                
                # 收集这个方法调用的函数
                callees[method_name] = self.call_graph['call_graph'].get(method_full_name, [])
        
        return {
            "class_name": class_name,
            "file_path": file_path,
            "methods": class_methods,
            "callers": callers,
            "callees": callees
        }
    
    def _get_call_graph_for_function(self, file_path, function_name):
        """获取函数的调用关系"""
        # 构建完整函数名
        full_function_name = f"{file_path}::{function_name}"
        
        # 获取调用这个函数的函数
        callers = self.call_graph['caller_graph'].get(full_function_name, [])
        
        # 获取这个函数调用的函数
        callees = self.call_graph['call_graph'].get(full_function_name, [])
        
        # 获取函数信息
        function_info = None
        for func_name, info in self.call_graph['function_info'].items():
            if info['file'] == file_path and info['name'] == function_name:
                function_info = info
                break
        
        return {
            "function_name": function_name,
            "file_path": file_path,
            "class_name": function_info.get('class', None) if function_info else None,
            "callers": [{"name": caller, "file": self.call_graph['function_info'].get(caller, {}).get('file', '未知')} 
                       for caller in callers if caller in self.call_graph['function_info']],
            "callees": [{"name": callee, "file": self.call_graph['function_info'].get(callee, {}).get('file', '未知')} 
                       for callee in callees if callee in self.call_graph['function_info']]
        }
